import asyncio
import click
from pathlib import Path
import os
import sys
from dotenv import load_dotenv
import subprocess
import signal

# [最终关键修复] 必须在导入任何 axiom_boot 框架模块之前，设置好 .env 文件的路径。
_env_file = Path(__file__).parent / '.env'
if not _env_file.exists():
    # 回退到 example 文件，确保至少有默认配置
    _env_file_example = Path(__file__).parent / 'env.example'
    if _env_file_example.exists():
        os.environ['ENV_FILE'] = str(_env_file_example.resolve())
        print(f"[WORKER Pre-Boot] .env 未找到，已加载 env.example: {os.environ['ENV_FILE']}")
    else:
        print(f"[WORKER Pre-Boot] 警告: 未找到 .env 或 env.example 文件。")
else:
    os.environ['ENV_FILE'] = str(_env_file.resolve())
    print(f"[WORKER Pre-Boot] 已加载 .env 文件: {os.environ['ENV_FILE']}")

    # [新增] 像 app.py 一样，强制性地预加载环境变量，确保配置被读取
    env_to_load = os.environ.get('ENV_FILE')
    if env_to_load:
        print(f"[WORKER Pre-Boot] 正在使用 load_dotenv 加载: {env_to_load}")
        load_dotenv(dotenv_path=env_to_load, override=True)
    else:
        print(f"[WORKER Pre-Boot] 未找到 ENV_FILE 环境变量，跳过 load_dotenv。")

    # 将 demo 项目自身的 src 目录添加到 sys.path
    _demo_src_path = Path(__file__).parent / 'src'
if str(_demo_src_path) not in sys.path:
    sys.path.insert(0, str(_demo_src_path))

# --- 环境准备完毕，现在可以安全导入框架 ---
from axiom_boot.di.application import AxiomApplication
from axiom_boot.task.arq_manager import ArqManager
from axiom_boot.core.exceptions import ServiceNotFoundException

# [关键修改] 不再使用 click.group，直接定义一个 command
@click.command()
@click.option('--queue', default=None, help='要监听的特定队列名称。')
@click.option('--workers', default=1, help='要启动的 Worker 进程数量。')
def run_worker_entry(queue: str | None, workers: int):
    """
    启动 ARQ 任务队列 Worker。
    这是此脚本的唯一入口点。
    """

    # 当需要启动多个 Worker 时，此脚本将作为进程管理器
    if workers > 1:
        click.echo(f"🚀 [WORKER-MANAGER] 正在启动 {workers} 个 Worker 进程...")
        processes = []
        # [关键修改] command 不再需要 'worker' 子命令
        command = [sys.executable, __file__]
        if queue:
            command.extend(['--queue', queue])

        for i in range(workers):
            proc = subprocess.Popen(command)
            processes.append(proc)
            click.echo(f"   - Worker #{i+1} 已启动 (PID: {proc.pid})")

        def signal_handler(sig, frame):
            click.echo("\n🚫 [WORKER-MANAGER] 收到终止信号，正在优雅地关闭所有 Worker 进程...")
            for p in processes:
                p.terminate()
            for p in processes:
                p.wait()
            sys.exit(0)

        signal.signal(signal.SIGINT, signal_handler)
        signal.signal(signal.SIGTERM, signal_handler)
        
        click.secho("--- Worker 管理器已启动。按 Ctrl+C 可关闭所有 Worker。 ---", fg="cyan")
        
        for p in processes:
            p.wait()
            
        click.echo("✅ [WORKER-MANAGER] 所有 Worker 进程已正常关闭。")
        return

    # --- 以下是 workers = 1 时的原始逻辑 ---
    click.echo("🚀 [WORKER] 正在启动 1 个 ARQ Worker...")

    async def _start_worker_instance():
        app = AxiomApplication(mode='worker')
        try:
            click.echo("   - [诊断] 即将刷新 DI 容器...")
            await app.context.refresh()
            click.echo("   - DI 容器 (Worker 模式) 初始化成功。")
            
            click.echo("   - 正在从容器中获取 ArqManager...")
            arq_manager = await app.context.get(ArqManager)
            click.echo(f"   - ArqManager 获取成功。准备监听队列: {queue or '默认'}")

            click.secho("--- ARQ Worker 即将启动 (run_worker)... ---", fg="green")
            await arq_manager.run_worker(app_context=app.context, queue_name=queue)
            click.secho("--- [诊断] ARQ Worker 的 run_worker 方法已返回，这意味着 Worker 已停止。 ---", fg="yellow")

        except ServiceNotFoundException:
            click.secho(f"   - 错误: 无法获取 ArqManager。", fg="red")
            click.secho(f"     请确保在 .env 或 env.example 文件中 'AXIOM_TASK__ENABLED=true' 已设置。", fg="red")
            return
        except Exception as e:
            click.secho(f"   - 发生未知错误: {e}", fg="red")
        finally:
            click.secho("--- [诊断] Worker 已停止或启动失败，正在进行最终的优雅关闭...", fg="yellow")
            if app.context and app.context._is_refreshed:
                await app.context.close()
            click.secho("--- [诊断] 应用上下文已关闭，Worker 协程即将退出。 ---", fg="yellow")

    try:
        asyncio.run(_start_worker_instance())
    finally:
        click.secho("--- [诊断] Worker 进程即将彻底退出。 ---", fg="red")

if __name__ == '__main__':
    # [关键修改] 直接调用唯一的命令函数
    run_worker_entry()
